"""
CTPRPCCommon

definitions of pseudo-constants, exceptions and functions used by
other CTPRPC components
"""

try:
	from hashlib import md5
	from hashlib import sha1 as sha
except:
	import md5
	import sha
import hmac
import select
import os
import struct

import CTPRPCDebug as debug
from CTPRPCPickle import CTPRPCSerializer

__revision__ = "2.00"

# global pseudo-constants
CTPRPC_VERSION_MAJOR	= 2
CTPRPC_VERSION_MINOR	= 0

# Encryption modes
CTP_ENCRYPT_NONE	= 0
CTP_ENCRYPT_DES3	= 1
CTP_ENCRYPT_AES		= 2

# Compression modes
CTP_COMPRESS_NONE	= 0
CTP_COMPRESS_BZ2	= 1
CTP_COMPRESS_ZLIB	= 2

# CTP-RPC commands
CMD_HEARTBEAT		= 0
CMD_INVOKE		= 1
CMD_HANDSHAKE_1		= 2
CMD_HANDSHAKE_2		= 3
CMD_LOGIN		= 4
CMD_SUBSCRIBE		= 5
CMD_UNSUBSCRIBE		= 6
CMD_READFILE		= 7

# Responses
RESP_OK			= 0
RESP_EXCEPTION		= 1
RESP_BADMETHOD		= 2
RESP_BADCOMMAND		= 3
RESP_HANDSHAKE_1	= 4
RESP_HANDSHAKE_2	= 5
RESP_BADLOGIN		= 6
RESP_BADEVENT		= 7

SERV_NOTIFY		= 100

# CTP Packet Flags
CTP_FLAG_COMPRESSED	= 0x01
CTP_FLAG_ENCRYPTED	= 0x02
CTP_FLAG_CONTINUED	= 0x04
CTP_FLAG_COMPRESS_DATA	= 0x08

# Other global settings
RSA_KEY_LENGTH		= 512
DEFAULT_KEY_SIZE	= 16
MAX_BLOB_SIZE		= 2	# MBytes
MAX_FILE_TRANSFER_SIZE	= 1	# MBytes
MAX_PACKET_SIZE		= 3	# MBytes


# *******************************************************
# Exception classes
class CTPRPCError ( Exception ): pass
class BadMethodError ( CTPRPCError ): pass
class BadEventError ( CTPRPCError ): pass
class InvalidResponseError ( CTPRPCError ): pass
class BadLoginOrPassword ( CTPRPCError ): pass
class BadProtocolVersion ( CTPRPCError ): pass


# MAIN

MAX_BLOB_SIZE *= 1048576
MAX_FILE_TRANSFER_SIZE *= 1048576
MAX_PACKET_SIZE *= 1048576


# build list of supported encryption/compression types
supported_encryption_types = {
				CTP_ENCRYPT_NONE:	[ '', None, 0 ],
				CTP_ENCRYPT_AES:	[ 'AES', None, 32 ],
				CTP_ENCRYPT_DES3:	[ 'DES3', None, 24 ],
			}

supported_compression_types = {
				CTP_COMPRESS_NONE:	[ '', None ],
				CTP_COMPRESS_BZ2:	[ 'bz2', None ],
				CTP_COMPRESS_ZLIB:	[ 'zlib', None ]
			}

#test if we can support encryption
try:
	import Crypto.PublicKey.RSA
	import Crypto.Util.randpool
	
	RSA = Crypto.PublicKey.RSA
	RANDPOOL = Crypto.Util.randpool.RandomPool()

except ImportError:
	supported_encryption_types = { CTP_ENCRYPT_NONE: [ '', None, 0 ] }

#test supported encryption algorithms
to_del = []
for type, mod in supported_encryption_types.iteritems():
	if not mod[0]: continue
	try:
		mod[1] = getattr ( __import__ ( 'Crypto.Cipher', globals(), locals(), [ mod[0] ] ), mod[0] )
		mod[1].key_size = mod[2]
	except ImportError:
		to_del.append ( type )
for d in to_del:
	del supported_encryption_types[d]

#test supported compression types
to_del = []
for type, mod in supported_compression_types.iteritems():
	if not mod[0]: continue
	try:
		mod[1] = __import__ ( mod[0] )
	except ImportError:
		to_del.append ( type )
for d in to_del:
	del supported_compression_types[d]

is_encryption_supported = len ( supported_encryption_types ) > 1
is_compression_supported = len ( supported_compression_types ) > 1


# *******************************************************
# CTP_PRF() - pseudo-random function
def CTP_PRF ( key, data, reqlen ):
	"""
	CTP_PRF ( key, data, reqlen ) -> str

	Pseudo-Random Function used internally by CTP-RPC server and client.
	"""
	
	md5_mac = hmac.new ( key, data, md5 ).digest()
	sha_mac = hmac.new ( key, data, sha ).digest()

	retval = ''
	while reqlen > 0:
		md5_new = hmac.new ( md5_mac + key, sha_mac + data, md5 ).digest()
		sha_new = hmac.new ( sha_mac + key, md5_mac + data, sha ).digest()
		retval += md5_new + sha_new
		md5_mac = md5_new
		sha_mac = sha_new
		reqlen -= 36
		
	return retval[:reqlen]
	

# *******************************************************
# read_bytes() - legge n bytes dal socket
def read_bytes ( sock, n ):
	"""
	read_bytes ( sock, n )

	Funzione di basso livello. Legge esattamente ``n`` byte da un socket.
	"""
	
	buf = ''
        timeout = sock.gettimeout()
        while n > 0:
                count = 0
                while count < timeout:
                        try: rl, wl, xl = select.select ( [ sock ], [], [ sock ], 1 )
                        except: return ''
                        if xl: return ''
                        elif rl: break
                        count += 1
                else:
                        return ''
                r_buf = sock.recv ( n )
                if not r_buf: return ''
                buf += r_buf
                n -= len ( r_buf )
        return buf


# *******************************************************
# read_object() - legge un oggetto da un socket
def read_object ( sock, session_id, hmac_key = '', compress_mod = None, encrypt_key = None ):
	"""
	read_object ( sock, session_id, hmac_key, compress_mod, encrypt_key ) -> obj or None

	Funzione CTP di basso livello per leggere un oggetto da un socket.

	Se la fase di handshake e' terminata effettua anche la decompressione/decifratura/verifica del MAC
	se sono state concordate.
	"""
	
	try:
		ret = ctp_sock_read ( sock, hmac_key, compress_mod, encrypt_key )
		if not ret: return None

		blob, flags = ret

		encrypted = flags & CTP_FLAG_ENCRYPTED
		compressed = flags & CTP_FLAG_COMPRESSED
		compress_data = flags & CTP_FLAG_COMPRESS_DATA

		ser = CTPRPCSerializer ()
		obj = ser.deserialize ( blob )

		#ANDREA: DEBUG
		#print "read_object: obj: %s" % obj
		
		files = ser.files [ : ]
		cur_file = None

		while len ( files ):
			
			ret = ctp_sock_read ( sock, hmac_key, compress_mod, encrypt_key )
			if not ret: return None

			buf = ret [ 0 ]
			
			fobj = ser.deserialize ( buf )
			
			if fobj [ 0 ] == 'FILE':
				if cur_file:
					cur_file.seek ( 0, 0 )
					del files [ 0 ]
				if not len ( files ): return None
				recv_fname = fobj [ 1 ]
				fname = files [ 0 ][ 0 ]
				if recv_fname != fname: return None
				cur_file = files [ 0 ][ 1 ]
			elif fobj [ 0 ] == 'DATA':
				if not cur_file: return None
				cur_file.write ( fobj [ 1 ] )
			elif fobj [ 0 ] == 'END':
				if cur_file:
					cur_file.seek ( 0, 0 )
					del files [ 0 ]
				if len ( files ): return None
			else:
				return None
		
		return ( obj, encrypted, compressed, compress_data )
	except:
		raise
		return None


# *******************************************************
# ctp_sock_read() - legge dati da un socket con il protocollo CTP
def ctp_sock_read ( sock, hmac_key, compress_mod, encrypt_key ):

	ser = CTPRPCSerializer ()

	blob = ''
	continued = 1
	while continued:

		flags = read_bytes ( sock, 1 )
		if not flags: return None
		flags = struct.unpack ( "B", flags ) [ 0 ]
		size = read_bytes ( sock, 4 )
		if not size: return None
		size = struct.unpack ( ">I", size ) [ 0 ]
		if size > MAX_PACKET_SIZE: return None
		buffer = read_bytes ( sock, size )
		if not buffer: return None

		encrypted	= flags & CTP_FLAG_ENCRYPTED
		compressed	= flags & CTP_FLAG_COMPRESSED
		continued	= flags & CTP_FLAG_CONTINUED

		#DECRYPTION
		if encrypted and encrypt_key:
			buffer = encrypt_key.decrypt ( buffer )
			size = struct.unpack ( ">I", buffer [ : 4 ] ) [ 0 ]
			if size + 4 > len ( buffer ): return None
			buffer = buffer [ 4 : 4 + size ]

		#DECOMPRESSION
		if compressed and compress_mod:
			buffer = compress_mod.decompress ( buffer )

		#MAC VERIFICATION
		if hmac_key:
			buffer, mac = ser.deserialize ( buffer )
			calc_mac = hmac.new ( hmac_key, buffer, md5 ).digest()
			calc_mac += hmac.new ( hmac_key, buffer, sha ).digest()
			if calc_mac != mac: return None

		blob += buffer

	return ( blob, flags )


# *******************************************************
# send_object() - invia un oggetto su un socket
def send_object ( sock, obj, session_id, hmac_key = '', compress_mod = None, encrypt_key = None, block_size = 1, compress_data = 1, relax_serialize = 0 ):
	"""
	send_object ( sock, obj, session_id, hmac_key, compress_mod, encrypt_key, block_size )

	Funzione CTP di basso livello per inviare un oggetto su un socket.

	Se la fase di handshake e' terminata effettua anche la compressione/cifratura/apposizione MAC
	se sono state concordate.
	"""

	ser = CTPRPCSerializer ( relax_serialize )
	blob = ser.serialize ( obj )

	files = ser.files [ : ]

	ctp_sock_write ( sock, blob, hmac_key, compress_mod, encrypt_key, block_size, compress_data )

	if not compress_data: compress_mod = None

	for f in files:

		#get file size
		#f.seek ( 0, 2 )
		#size = f.tell()
		#f.seek ( 0, 0 )

		buf = ser.serialize ( ( 'FILE', f.name ) )
		ctp_sock_write ( sock, buf, hmac_key, compress_mod, encrypt_key, block_size, compress_data )

		while 1:
			buf = f.read ( MAX_FILE_TRANSFER_SIZE )
			if not len ( buf ): break

			buf = ser.serialize ( ( 'DATA', buf ) )
			ctp_sock_write ( sock, buf, hmac_key, compress_mod, encrypt_key, block_size, compress_data )
			
			#if len ( buf ) < MAX_FILE_TRANSFER_SIZE: break
	
	if len ( files ):
		buf = ser.serialize ( ( 'END', ) )
		ctp_sock_write ( sock, buf, hmac_key, compress_mod, encrypt_key, block_size, compress_data )


# *******************************************************
# ctp_sock_write() - scrive dati su un socket con il protocollo CTP
def ctp_sock_write ( sock, blob, hmac_key, compress_mod, encrypt_key, block_size, compress_data = 1 ):

	ser = CTPRPCSerializer ()

	size = len ( blob )
	
	cur_ptr = 0
	while size > 0:

		flags = 0
		if compress_data: flags |= CTP_FLAG_COMPRESS_DATA
		
		if size > MAX_BLOB_SIZE:
			size -= MAX_BLOB_SIZE
			to_send = MAX_BLOB_SIZE
			flags |= CTP_FLAG_CONTINUED
		else:
			to_send = size
			size = 0

		obj_str = blob [ cur_ptr : cur_ptr + to_send ]
		cur_ptr += to_send

		#MAC
		if hmac_key:
			mac = hmac.new ( hmac_key, obj_str, md5 ).digest ()
			mac += hmac.new ( hmac_key, obj_str, sha ).digest ()
			obj = ( obj_str, mac )
			obj_str = ser.serialize ( obj )

		#COMPRESSION
		if compress_mod:
			comp_obj_str = compress_mod.compress ( obj_str )
			if len ( comp_obj_str ) < len ( obj_str ):
				obj_str = comp_obj_str
				flags |= CTP_FLAG_COMPRESSED

		#ENCRYPTION
		if encrypt_key:
			obj_str = struct.pack ( ">I", len ( obj_str ) ) + obj_str
			padlen = block_size - ( len ( obj_str ) % block_size )
			obj_str += "J" * padlen
			obj_str = encrypt_key.encrypt ( obj_str )
			flags |= CTP_FLAG_ENCRYPTED

		sock.sendall ( struct.pack ( ">BI", flags, len ( obj_str ) ) )
		sock.sendall ( obj_str )


def send_welcome ( sock ):
	wm = "GB" + struct.pack ( "BB", CTPRPC_VERSION_MAJOR, CTPRPC_VERSION_MINOR )
	sock.sendall ( wm )


def get_welcome ( sock ):
	wm = read_bytes ( sock, 4 )
	if not wm: return None
	if wm [ 0 ] != 'G': return None
	if wm [ 1 ] != 'B': return None
	return struct.unpack ( "BB", wm [ 2 : ] )


def build_pkcs_block_02 ( data, modlen ):
	l = len ( data )
	modlen /= 8
	if l >= modlen - 3:
		raise Exception, "cannot build RSA block 02 (data length >= modulus length - 3)"

	padlen = modlen - l - 3
	rnd = ""
	for i in xrange ( padlen ):
		b = RANDPOOL.get_bytes ( 1 )
		if b == '\x00': b = '\x69'
		rnd += b

	return "\x00\x02" + rnd + "\x00" + data


def read_pkcs_block_02 ( data, modlen ):
	l = len ( data )
	modlen /= 8

	if l == modlen:
		if data [ 0 ] != "\x00": return None
		data = data [ 1 : ]
	elif l != modlen - 1: return None
	
	if data [ 0 ] != "\x02": return None

	for i, b in enumerate ( data [ 1 : ] ):
		if b == "\x00": break
	else:
		return None

	return data [ 2 + i : ]

